// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <string>

#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/array_shuffle.h>
#include <pollux/functions/lib/repeat.h>
#include <pollux/functions/lib/slice.h>
#include <pollux/functions/prestosql/array_constructor.h>
#include <pollux/functions/prestosql/array_functions.h>
#include <pollux/functions/prestosql/array_sort.h>
#include <pollux/functions/prestosql/width_bucket_array.h>
#include <pollux/functions/prestosql/types/json_registration.h>

namespace kumo::pollux::functions {
    extern void registerArrayConcatFunctions(const std::string &prefix);

    extern void registerArrayNGramsFunctions(const std::string &prefix);

    template<typename T>
    inline void registerArrayMinMaxFunctions(const std::string &prefix) {
        register_function<ArrayMinFunction, T, Array<T> >({prefix + "array_min"});
        register_function<ArrayMaxFunction, T, Array<T> >({prefix + "array_max"});
    }

    template<typename T>
    inline void registerArrayJoinFunctions(const std::string &prefix) {
        register_function<
            ParameterBinder<ArrayJoinFunction, T>,
            Varchar,
            Array<T>,
            Varchar>({prefix + "array_join"});

        register_function<
            ParameterBinder<ArrayJoinFunction, T>,
            Varchar,
            Array<T>,
            Varchar,
            Varchar>({prefix + "array_join"});
    }

    template<typename T>
    inline void registerArrayCombinationsFunctions(const std::string &prefix) {
        register_function<
            ParameterBinder<CombinationsFunction, T>,
            Array<Array<T> >,
            Array<T>,
            int32_t>({prefix + "combinations"});
    }

    template<typename T>
    inline void registerArrayCumSumFunction(const std::string &prefix) {
        register_function<ParameterBinder<ArrayCumSumFunction, T>, Array<T>, Array<T> >(
            {prefix + "array_cum_sum"});
    }

    template<typename T>
    inline void registerArrayHasDuplicatesFunctions(const std::string &prefix) {
        register_function<
            ParameterBinder<ArrayHasDuplicatesFunction, T>,
            bool,
            Array<T> >({prefix + "array_has_duplicates"});
    }

    template<typename T>
    inline void registerArrayFrequencyFunctions(const std::string &prefix) {
        register_function<
            ParameterBinder<ArrayFrequencyFunction, T>,
            Map<T, int>,
            Array<T> >({prefix + "array_frequency"});
    }

    template<typename T>
    inline void registerArrayNormalizeFunctions(const std::string &prefix) {
        register_function<
            ParameterBinder<ArrayNormalizeFunction, T>,
            Array<T>,
            Array<T>,
            T>({prefix + "array_normalize"});
    }

    template<typename T>
    inline void registerArrayTrimFunctions(const std::string &prefix) {
        register_function<ArrayTrimFunction, Array<T>, Array<T>, int64_t>(
            {prefix + "trim_array"});
    }

    template<typename T>
    inline void registerArrayTopNFunction(const std::string &prefix) {
        register_function<ArrayTopNFunction, Array<T>, Array<T>, int32_t>(
            {prefix + "array_top_n"});
    }

    template<typename T>
    inline void registerArrayRemoveNullFunctions(const std::string &prefix) {
        register_function<ArrayRemoveNullFunction, Array<T>, Array<T> >(
            {prefix + "remove_nulls"});
    }

    template<typename T>
    inline void registerArrayUnionFunctions(const std::string &prefix) {
        register_function<ArrayUnionFunction, Array<T>, Array<T>, Array<T> >(
            {prefix + "array_union"});
    }

    template<typename T>
    inline void registerArrayRemoveFunctions(const std::string &prefix) {
        register_function<ArrayRemoveFunction, Array<T>, Array<T>, T>(
            {prefix + "array_remove"});
    }

    void registerInternalArrayFunctions() {
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_$internal$canonicalize, "$internal$canonicalize");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_$internal$contains, "$internal$contains");
    }

    void registerArrayFunctions(const std::string &prefix) {
        registerJsonType();
        registerArrayConstructor(prefix + "array_constructor");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_all_match, prefix + "all_match");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_any_match, prefix + "any_match");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_none_match, prefix + "none_match");

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_find_first, prefix + "find_first");
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_find_first_index, prefix + "find_first_index");

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_distinct, prefix + "array_distinct");
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_array_duplicates, prefix + "array_duplicates");
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_array_intersect, prefix + "array_intersect");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_contains, prefix + "contains");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_except, prefix + "array_except");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_arrays_overlap, prefix + "arrays_overlap");
        registerBigintSliceFunction(prefix);
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_zip, prefix + "zip");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_zip_with, prefix + "zip_with");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_position, prefix + "array_position");
        exec::registerStatefulVectorFunction(
            prefix + "shuffle",
            arrayShuffleSignatures(),
            makeArrayShuffle,
            getMetadataForArrayShuffle());

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_sort, prefix + "array_sort");
        POLLUX_REGISTER_VECTOR_FUNCTION(
            udf_array_sort_desc, prefix + "array_sort_desc");

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_max_by, prefix + "array_max_by");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_min_by, prefix + "array_min_by");

        exec::registerExpressionRewrite([prefix](const auto &expr) {
            return rewriteArraySortCall(prefix, expr);
        });

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_array_sum, prefix + "array_sum");
        exec::registerStatefulVectorFunction(
            prefix + "repeat", repeatSignatures(), makeRepeat, repeatMetadata());
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_sequence, prefix + "sequence");

        exec::registerStatefulVectorFunction(
            prefix + "width_bucket",
            widthBucketArraySignature(),
            makeWidthBucketArray);

        registerArrayMinMaxFunctions<int8_t>(prefix);
        registerArrayMinMaxFunctions<int16_t>(prefix);
        registerArrayMinMaxFunctions<int32_t>(prefix);
        registerArrayMinMaxFunctions<int64_t>(prefix);
        registerArrayMinMaxFunctions<int128_t>(prefix);
        registerArrayMinMaxFunctions<float>(prefix);
        registerArrayMinMaxFunctions<double>(prefix);
        registerArrayMinMaxFunctions<bool>(prefix);
        registerArrayMinMaxFunctions<Varchar>(prefix);
        registerArrayMinMaxFunctions<Timestamp>(prefix);
        registerArrayMinMaxFunctions<Date>(prefix);
        registerArrayMinMaxFunctions<Orderable<T1> >(prefix);

        registerArrayJoinFunctions<int8_t>(prefix);
        registerArrayJoinFunctions<int16_t>(prefix);
        registerArrayJoinFunctions<int32_t>(prefix);
        registerArrayJoinFunctions<int64_t>(prefix);
        registerArrayJoinFunctions<int128_t>(prefix);
        registerArrayJoinFunctions<float>(prefix);
        registerArrayJoinFunctions<double>(prefix);
        registerArrayJoinFunctions<bool>(prefix);
        registerArrayJoinFunctions<Varchar>(prefix);
        registerArrayJoinFunctions<Timestamp>(prefix);
        registerArrayJoinFunctions<Date>(prefix);
        registerArrayJoinFunctions<Json>(prefix);
        registerArrayJoinFunctions<UnknownValue>(prefix);

        register_function<ArrayAverageFunction, double, Array<double> >(
            {prefix + "array_average"});

        registerArrayConcatFunctions(prefix);
        registerArrayNGramsFunctions(prefix);

        register_function<
            ArrayFlattenFunction,
            Array<Generic<T1> >,
            Array<Array<Generic<T1> > > >({prefix + "flatten"});

        registerArrayRemoveFunctions<int8_t>(prefix);
        registerArrayRemoveFunctions<int16_t>(prefix);
        registerArrayRemoveFunctions<int32_t>(prefix);
        registerArrayRemoveFunctions<int64_t>(prefix);
        registerArrayRemoveFunctions<int128_t>(prefix);
        registerArrayRemoveFunctions<float>(prefix);
        registerArrayRemoveFunctions<double>(prefix);
        registerArrayRemoveFunctions<bool>(prefix);
        registerArrayRemoveFunctions<Timestamp>(prefix);
        registerArrayRemoveFunctions<Date>(prefix);
        registerArrayRemoveFunctions<Varbinary>(prefix);
        registerArrayRemoveFunctions<Generic<T1> >(prefix);
        register_function<
            ArrayRemoveFunctionString,
            Array<Varchar>,
            Array<Varchar>,
            Varchar>({prefix + "array_remove"});

        registerArrayTrimFunctions<int8_t>(prefix);
        registerArrayTrimFunctions<int16_t>(prefix);
        registerArrayTrimFunctions<int32_t>(prefix);
        registerArrayTrimFunctions<int64_t>(prefix);
        registerArrayTrimFunctions<int128_t>(prefix);
        registerArrayTrimFunctions<float>(prefix);
        registerArrayTrimFunctions<double>(prefix);
        registerArrayTrimFunctions<bool>(prefix);
        registerArrayTrimFunctions<Timestamp>(prefix);
        registerArrayTrimFunctions<Date>(prefix);
        registerArrayTrimFunctions<Varbinary>(prefix);
        registerArrayTrimFunctions<Generic<T1> >(prefix);
        register_function<
            ArrayTrimFunctionString,
            Array<Varchar>,
            Array<Varchar>,
            int64_t>({prefix + "trim_array"});

        registerArrayTopNFunction<int8_t>(prefix);
        registerArrayTopNFunction<int16_t>(prefix);
        registerArrayTopNFunction<int32_t>(prefix);
        registerArrayTopNFunction<int64_t>(prefix);
        registerArrayTopNFunction<int128_t>(prefix);
        registerArrayTopNFunction<float>(prefix);
        registerArrayTopNFunction<double>(prefix);
        registerArrayTopNFunction<Varchar>(prefix);
        registerArrayTopNFunction<Timestamp>(prefix);
        registerArrayTopNFunction<Date>(prefix);
        registerArrayTopNFunction<Varbinary>(prefix);
        registerArrayTopNFunction<Orderable<T1> >(prefix);

        registerArrayRemoveNullFunctions<int8_t>(prefix);
        registerArrayRemoveNullFunctions<int16_t>(prefix);
        registerArrayRemoveNullFunctions<int32_t>(prefix);
        registerArrayRemoveNullFunctions<int64_t>(prefix);
        registerArrayRemoveNullFunctions<int128_t>(prefix);
        registerArrayRemoveNullFunctions<float>(prefix);
        registerArrayRemoveNullFunctions<double>(prefix);
        registerArrayRemoveNullFunctions<bool>(prefix);
        registerArrayRemoveNullFunctions<Timestamp>(prefix);
        registerArrayRemoveNullFunctions<Date>(prefix);
        registerArrayRemoveNullFunctions<Varbinary>(prefix);
        registerArrayRemoveNullFunctions<Generic<T1> >(prefix);
        register_function<
            ArrayRemoveNullFunctionString,
            Array<Varchar>,
            Array<Varchar> >({prefix + "remove_nulls"});

        registerArrayUnionFunctions<int8_t>(prefix);
        registerArrayUnionFunctions<int16_t>(prefix);
        registerArrayUnionFunctions<int32_t>(prefix);
        registerArrayUnionFunctions<int64_t>(prefix);
        registerArrayUnionFunctions<int128_t>(prefix);
        registerArrayUnionFunctions<float>(prefix);
        registerArrayUnionFunctions<double>(prefix);
        registerArrayUnionFunctions<bool>(prefix);
        registerArrayUnionFunctions<Timestamp>(prefix);
        registerArrayUnionFunctions<Date>(prefix);
        registerArrayUnionFunctions<Varbinary>(prefix);
        registerArrayUnionFunctions<Generic<T1> >(prefix);

        registerArrayCombinationsFunctions<int8_t>(prefix);
        registerArrayCombinationsFunctions<int16_t>(prefix);
        registerArrayCombinationsFunctions<int32_t>(prefix);
        registerArrayCombinationsFunctions<int64_t>(prefix);
        registerArrayCombinationsFunctions<int128_t>(prefix);
        registerArrayCombinationsFunctions<float>(prefix);
        registerArrayCombinationsFunctions<double>(prefix);
        registerArrayCombinationsFunctions<bool>(prefix);
        registerArrayCombinationsFunctions<Varchar>(prefix);
        registerArrayCombinationsFunctions<Timestamp>(prefix);
        registerArrayCombinationsFunctions<Date>(prefix);
        registerArrayCombinationsFunctions<Generic<T1> >(prefix);

        registerArrayCumSumFunction<int8_t>(prefix);
        registerArrayCumSumFunction<int16_t>(prefix);
        registerArrayCumSumFunction<int32_t>(prefix);
        registerArrayCumSumFunction<int64_t>(prefix);
        registerArrayCumSumFunction<int128_t>(prefix);
        registerArrayCumSumFunction<float>(prefix);
        registerArrayCumSumFunction<double>(prefix);

        registerArrayHasDuplicatesFunctions<int8_t>(prefix);
        registerArrayHasDuplicatesFunctions<int16_t>(prefix);
        registerArrayHasDuplicatesFunctions<int32_t>(prefix);
        registerArrayHasDuplicatesFunctions<int64_t>(prefix);
        registerArrayHasDuplicatesFunctions<int128_t>(prefix);
        registerArrayHasDuplicatesFunctions<Varchar>(prefix);
        registerArrayHasDuplicatesFunctions<Json>(prefix);

        registerArrayFrequencyFunctions<bool>(prefix);
        registerArrayFrequencyFunctions<int8_t>(prefix);
        registerArrayFrequencyFunctions<int16_t>(prefix);
        registerArrayFrequencyFunctions<int32_t>(prefix);
        registerArrayFrequencyFunctions<int64_t>(prefix);
        registerArrayFrequencyFunctions<int128_t>(prefix);
        registerArrayFrequencyFunctions<float>(prefix);
        registerArrayFrequencyFunctions<double>(prefix);
        registerArrayFrequencyFunctions<Timestamp>(prefix);
        registerArrayFrequencyFunctions<Date>(prefix);
        registerArrayFrequencyFunctions<Varchar>(prefix);

        registerArrayNormalizeFunctions<int8_t>(prefix);
        registerArrayNormalizeFunctions<int16_t>(prefix);
        registerArrayNormalizeFunctions<int32_t>(prefix);
        registerArrayNormalizeFunctions<int64_t>(prefix);
        registerArrayNormalizeFunctions<float>(prefix);
        registerArrayNormalizeFunctions<double>(prefix);
    }
} // namespace kumo::pollux::functions
